#################################
##' Code for:
##' Title: "Effective climate clubs require ambition, leverage, and insulation"
##' Author: Sam S. Rowan
##' Contact: sam.rowan [at] concordia.ca
##' Date: 2023-10-19
#################################
##' This script:
##' Distribution of membership across simulations/draws
#################################


#### Initialize environment ####

## Packages
library(tidyverse)
library(truncnorm)
library(knitr)

## Working directory
# setwd("/Users/srowa/Dropbox/Projects/climate_clubs_trade")
# -- Set working directory to replication folder


## Define the logistic function
logistic_function <- function(x, k, a) {
  return(1 / (1 + exp(-k * (x - a))))
}

## Plot settings for later
theme_clubs <-  theme(legend.position = "bottom", 
                      panel.background = element_rect(fill=NA),
                      panel.grid.major.x = element_blank(),
                      panel.grid.major.y = element_blank(),
                      panel.grid.minor.y =  element_blank(),
                      panel.grid.minor.x =  element_blank(),
                      panel.ontop = FALSE,
                      axis.line = element_line(linewidth = 1, colour = "black"),
                      axis.text = element_text(size = 16),
                      legend.text = element_text(size = 12),
                      legend.title = element_text(size = 16),
                      axis.title = element_text(size = 16))


#### Define climate model parameters ####

## Define number of states
num_states <- 150

## Draw their climate ideal points
set.seed(202109)
climate_ideal <- cbind.data.frame(
  climate_ideal = runif(n = num_states, min = 0, max = 1)) %>% 
  arrange(climate_ideal) %>% 
  mutate(id = seq(from = 1, to = num_states, by = 1))
# climate_ideal
#' Leave climate ideal points the same
#' But we'll draw trade ties 500 times

####' ---- NOTE: Code block below runs for about 5'
####' ---- The simulation output is saved as 'output/members_baseline_theta75.csv'
####' ---- Skip running the simulation and load this in directly at around line 300
####' ---- And build figures from there

#### Simulate trade ties 500 times, analyze for fixed theta ####

## Create loop counters
set.seed(202109)
seeds_vector <- as.integer(runif(n = 500, min = 1, max = 1e6))
theta_vector <- 0.75

## Create objects to hold outputs
# Count membership for each value of theta and each draw of trade ties
count_members_independent <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
count_members_positive <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))
count_members_negative <- matrix(NA, nrow = length(theta_vector), ncol = length(seeds_vector))


## Simulate dyadic trade flows 500 times to get frequency distribution
start_time <- Sys.time() # Takes 5 minutes on 2021 iMac
for (seed in seeds_vector) {
  # Declare seed
  set.seed(seed)
  theta_df <-  cbind.data.frame(
    # Add IDs
    id_a = rep(climate_ideal$id, times = num_states),
    id_b = rep(climate_ideal$id, each = num_states),
    # Bind climate ideal points into dyadic setup
    climate_ideal_a = rep(climate_ideal$climate_ideal, times = num_states),
    climate_ideal_b = rep(climate_ideal$climate_ideal, each = num_states)) %>% 
    # Measure dyadic ideal point distance
    mutate(ideal_point_distance = abs(climate_ideal_a - climate_ideal_b)) %>% 
    # --** Random component **--
    # Draw trade data based on ideal point distance
    mutate(trade_independent = rtruncnorm(n = num_states^2,
                                          # a = 0, b = 1,
                                          mean = 0,
                                          sd = 0.5),
           trade_positive = rtruncnorm(n = num_states^2,
                                       # a = 0, b = 1,
                                       mean = (1-ideal_point_distance), # negative of difference, implies trade more when ideal points closer
                                       sd = 0.5),
           trade_negative = rtruncnorm(n = num_states^2,
                                       # a = 0, b = 1,
                                       mean = ideal_point_distance, # positive of difference, implies trade less when ideal points closer
                                       sd = 0.5)) %>% 
    # Use logistic function to bound trade ties
    mutate_at(vars(trade_independent:trade_negative),
              ~logistic_function(x = ., k = 10, a = 0.75)) %>% 
    # Remove data for faux dyads -- where id_a == id_b
    mutate_at(vars(ideal_point_distance,
                   trade_independent, trade_positive, trade_negative),
              ~replace(., id_a == id_b, NA)) %>% 
    # Get each state's total trade
    group_by(id_a) %>% 
    mutate(sum_trade_independent = sum(trade_independent, na.rm = T),
           sum_trade_positive = sum(trade_positive, na.rm = T),
           sum_trade_negative = sum(trade_negative, na.rm = T)) %>% 
    # Standardize
    rowwise() %>% 
    mutate(share_dyadic_independent = trade_independent / sum_trade_independent,
           share_dyadic_positive = trade_positive / sum_trade_positive,
           share_dyadic_negative = trade_negative / sum_trade_negative) %>% 
    select(id_a, id_b, climate_ideal_a, climate_ideal_b, ideal_point_distance,
           share_dyadic_independent, share_dyadic_positive, share_dyadic_negative) %>% 
    # Create a club's common climate goal
    mutate(theta = theta_vector,
           core_member_a = ifelse(climate_ideal_a >= theta, 1, 0),
           core_member_b = ifelse(climate_ideal_b >= theta, 1, 0)) %>% 
    # Flag trade to club members
    mutate(trade_to_club_independent = ifelse(core_member_b == 1, 
                                              share_dyadic_independent, NA),
           trade_to_club_positive = ifelse(core_member_b == 1, 
                                           share_dyadic_positive, NA),
           trade_to_club_negative = ifelse(core_member_b == 1, 
                                           share_dyadic_negative, NA)) %>% 
    # Get share of trade to all club members
    group_by(id_a) %>% 
    mutate(share_trade_club_independent = sum(trade_to_club_independent, na.rm = T),
           share_trade_club_positive = sum(trade_to_club_positive, na.rm = T),
           share_trade_club_negative = sum(trade_to_club_negative, na.rm = T)) %>% 
    ungroup() %>% 
    select(-starts_with("trade_to_club")) %>% 
    # Collapse the data to just _a's perspective, to calculate utilities
    select(id_a, climate_ideal_a, theta, core_member_a,
           starts_with("share_trade_club")) %>% 
    distinct() %>% # rows = 150
    # Calculate utilities to membership
    mutate(beta = 1) %>% # Define weights for climate versus trade
    rowwise() %>% 
    # Get most constrained core club member
    mutate(min_trade_exposure_independent = ifelse(
      core_member_a == 1, min(share_trade_club_independent), NA),
      min_trade_exposure_positive = ifelse(
        core_member_a == 1, min(share_trade_club_positive ), NA),
      min_trade_exposure_negative = ifelse(
        core_member_a == 1, min(share_trade_club_negative), NA)) %>% 
    ungroup() %>% 
    # Define that constraint as lambda
    mutate(lambda_independent = min(min_trade_exposure_independent, na.rm = T),
           lambda_positive = min(min_trade_exposure_positive, na.rm = T),
           lambda_negative = min(min_trade_exposure_negative, na.rm = T)) %>% 
    # Calculate utilities for the independent world
    # Utility to join
    mutate(utility_join_independent = (-abs(climate_ideal_a - theta)) + # Climate
             (1 - share_trade_club_independent) * (1 - lambda_independent)*beta + # Non-club trade
             (share_trade_club_independent*beta) - # Club trade
             (theta - lambda_independent)) %>%  # Aversion
    # Utility to abstain
    mutate(utility_abstain_independent = (-climate_ideal_a) + # Climate
             (1-share_trade_club_independent)*beta + # Non-club trade
             (share_trade_club_independent*(1-lambda_independent)*beta) + # Club trade
             (0)) %>%  # Aversion
    # Net, membership
    mutate(net_utility_independent = utility_join_independent - utility_abstain_independent,
           membership_independent = ifelse(net_utility_independent >= 0, 1, 0)) %>% 
    # Calculate utilities for the positive world
    # Utility to join
    mutate(utility_join_positive = (-abs(climate_ideal_a - theta)) + # Climate
             (1 - share_trade_club_positive) * (1 - lambda_positive)*beta + # Non-club trade
             (share_trade_club_positive*beta) - # Club trade
             (theta - lambda_positive)) %>%  # Aversion
    # Utility to abstain
    mutate(utility_abstain_positive = (-climate_ideal_a) + # Climate
             (1-share_trade_club_positive)*beta + # Non-club trade
             (share_trade_club_positive*(1-lambda_positive)*beta) + # Club trade
             (0)) %>%  # Aversion
    # Net, membership
    mutate(net_utility_positive = utility_join_positive - utility_abstain_positive,
           membership_positive = ifelse(net_utility_positive >= 0, 1, 0)) %>% 
    # Calculate utilities for the negative world
    # Utility to join
    mutate(utility_join_negative = (-abs(climate_ideal_a - theta)) + # Climate
             (1 - share_trade_club_negative) * (1 - lambda_negative)*beta + # Non-club trade
             (share_trade_club_negative*beta) - # Club trade
             (theta - lambda_negative)) %>%  # Aversion
    # Utility to abstain
    mutate(utility_abstain_negative = (-climate_ideal_a) + # Climate
             (1-share_trade_club_negative)*beta + # Non-club trade
             (share_trade_club_negative*(1-lambda_negative)*beta) + # Club trade
             (0)) %>%  # Aversion
    # Net, membership
    mutate(net_utility_negative = utility_join_negative - utility_abstain_negative,
           membership_negative = ifelse(net_utility_negative >= 0, 1, 0))
  
  # Save the number of countries in the club
  count_members_independent[which(seed == seeds_vector)] <- nrow(
    filter(theta_df, membership_independent == 1))
  count_members_positive[which(seed == seeds_vector)] <- nrow(
    filter(theta_df, membership_positive == 1))
  count_members_negative[which(seed == seeds_vector)] <- nrow(
    filter(theta_df, membership_negative == 1))
}
end_time <- Sys.time()
end_time - start_time


#### Summarize outputs ####

## Tidy outputs from simulation
members_baseline <- cbind.data.frame(independent = count_members_independent[1,],
                                     positive = count_members_positive[1,],
                                     negative = count_members_negative[1,]) %>% 
  as.data.frame() %>% 
  mutate(run = "Baseline")

# # Write to disk -- redundant, not run
# write_csv(x = members_baseline,
#           file = "output/members_baseline_theta75.csv")

## Load from disk -- redundant if ran simulation
members_baseline <- read_csv("output/members_baseline_theta75.csv")


## Table
members_baseline %>% 
  pivot_longer(names_to = "model", 
               values_to = "members",
               independent:negative) %>% 
  group_by(model) %>% 
  summarize(median = median(members)) %>% 
  mutate(prop.median = median/150)

  
## Plot
members_at_theta75 <- members_baseline %>% 
  pivot_longer(names_to = "model", 
               values_to = "members",
               independent:negative) %>% 
  group_by(model, members) %>% 
  summarize(n = n()) %>% 
  mutate(model = case_when(model == "independent" ~ "Independent",
                           model == "positive" ~ "Positive",
                           model == "negative" ~ "Negative")) %>% 
  ggplot(., aes(x = members, y = n, fill = model)) +
  # geom_col(position = "dodge") +
  geom_col(position = position_dodge2(width = 0.9, preserve = "single")) +
  scale_fill_manual(values = c("#7DCEA0", "#229954", "#556B2F")) +
  scale_x_continuous(breaks = c(38, 40, 45, 50, 55, 58), 
                     labels = c(38, 40, 45, 50, 55, 58)) +
  labs(x = expression(Club~members~at~theta[j]==0.75),
       y = "Number of simulations",
       fill = "Model") +
  theme_clubs
members_at_theta75
ggsave(plot = members_at_theta75, filename = "text/members_at_theta75.pdf", units = 'in', height = 6, width = 8)









